#!/usr/bin/env python3

"""Fast Fourier Transform for functions on [0,1]
"""

import numpy as np


def alt_sequences(k, start=1):
    a, r = divmod(k, 2)
    seq = np.tile([1, -1], a)
    if r:
        seq = np.append(seq, 1)
    return -seq if start ==-1 else seq


def fc(f, Nt=100, krange=None, padding='c', tstart=0):
    """Fourier coefficents
    
    Ff_k = int_0^1 f(x)exp{-2pi i kx} dx
    calcualat it with Trapzoid integral formula.

    Caution: krange2-krange1 <= Nt-1
    
    Arguments:
        f {function on [0,1] | array} -- function on [0,1] as a 1D signal
    
    Keyword Arguments:
        Nt + 1 {number} -- number of samples on time (space)(default: {100})
        krange {tuple} -- range of index (in Z) (default: {None})
        padding {str} -- Used in feature
    
    Returns:
        F, w : array[Nw]
            """

    if callable(f):
        t = np.linspace(0, 1, Nt + 1)
        f = f(t)
    else:
        if Nt is None:
            Nt = f.shape[0]
        else:
            f = f[:Nt]
        t = np.linspace(0, 1, Nt + 1)
        if padding == 'c':
            padding_value = f[-1]

    if krange is None:
        k1 = - L // 2
        k2 = L + k1 - 1
        krange = k1, k2
    else:
        k1, k2 = krange

    assert k2 - k1 <= Nt, 'The nubmer of coefficients must be less than Nt!'

    c = np.fft.fft(f * np.exp(-2j*np.pi*k1*t[:-1]))
    c = (c + (padding_value - f[0])/2) / Nt
    if tstart=='s':
        return c[:k2-k1] * alt_sequences(k2-k1, -k1%2)
    elif tstart == 0:
        return c[:k2-k1]
    else:
        return c[:k2-k1] * np.exp(-2j*np.pi*tstart*np.arange(k1, k2))


def fr(c, krange=None, Nt=100, padding='c'):
    """Fourier reconstruction
    f(t)=sum_nc_nexp(2pi jnt) : L^2[0,1]
    Caution: krange2-krange1 <= Nt
    """

    L = c.shape[0]
    if krange is None:
        k1 = - L // 2
        k2 = L + k1 - 1
        krange = k1, k2
    else:
       k1, k2 = krange
    t = np.linspace(0, 1, Nt+1)
    d = Nt - L
    if d >= 0:
        c = np.concatenate([c, np.zeros((d,)+c.shape[1:])], axis=0)
    f = np.fft.ifft(c) * np.exp(2j*np.pi*k1*t[:-1]) * Nt
    return f, t[:-1]


def fapprox(f, Nt=100, krange=(-30,30), threshold=0.1, padding='c'):

    if callable(f):
        t = np.linspace(0, 1, Nt+1)
        f = f(t)
    else:
        if Nt is None:
            Nt = f.shape[0]
        else:
            f = f[:Nt]
    c = fc(f, Nt, krange, padding=padding)
    cm = np.max(np.abs(c))
    c = np.array([ci if np.abs(ci) > threshold * cm else 0 for ci in c])
    return fr(c, krange, Nt)


def fc2(f, Nt=(100, 100), k1range=(-10, 10), k2range=(-10,10), padding='c'):
    """Calculate Fourier coefficents for 2dim functions
        f : function on [0,1]^2 | array;
        Nt+1: number of points of space
        krange : range of index (Z^2);
    Output:
        F, w : array[Nw]
    Call: fc
    """

    if callable(f):
        x = np.linspace(0, 1, Nt[0]+1)
        y = np.linspace(0, 1, Nt[1]+1)
        f = f(x,y)
    else:
        if Nt is None:
            Nt = f.shape[0], f.shape[1]
        else:
            f = f[:Nt[0], :Nt[1]]
        x = np.linspace(0, 1, Nt[0]+1)
        y = np.linspace(0, 1, Nt[1]+1)
    g = np.array([fc(row, Nt[1], k2range, padding=padding) for row in f])
    F = np.array([fc(col, Nt[0], k1range, padding=padding) for col in g.transpose()])
    return F.transpose()


def fr2(c, k1range=None, k2range=None, Nt=(100, 100), padding='c'):
    """ Fourier reconstruction for 2-dim functions
    """

    L1, L2 = c.shape[0:2]
    if k1range is None:
        k1 = - L // 2
        k2 = L + k1 - 1
        k1range = k1, k2
    if k2range is None:
        k1 = - L2 // 2
        k2 = L2 +k1 - 1
        k2range = k1, k2

    x = np.linspace(0, 1, Nt[0]+1)
    y = np.linspace(0, 1, Nt[1]+1)
    F = np.array([fr(row, k2range, Nt[1], padding=padding)[0] for row in c])
    F = np.array([fr(col, k1range, Nt[0], padding=padding)[0] for col in F.transpose()])
    return F.transpose(), x[:-1], y[:-1]

def fapprox2(f, Nt=(100, 100), k1range=(-30,30), k2range=(-30,30), threshold=0.0, padding='c'):

    if callable(f):
        x = np.linspace(0, 1, Nt[0])
        y = np.linspace(0, 1, Nt[1])
        f = f(x, y)
    else:
        if Nt is None:
            Nt = f.shape[0], f.shape[1]
        else:
            f = f[:Nt[0]+1,:Nt[1]+1]
    c = fc2(f, Nt, k1range, k2range, padding=padding)
    cm = np.max(np.abs(c))
    c = np.array([[ci if np.abs(ci) > threshold * cm else 0 for ci in row] for row in c])
    return fr2(c, k1range, k2range, Nt, padding=padding)


class Signal:
    pass

 
class SignalProcessor:
    def __init__(self, signal):
        self.signal = signal


class SignalProcessor1D(SignalProcessor):
    # Image Processor with Fourier Analysis
    # signal: function on [0,1] or vector
    def __init__(self, signal, krange=(-50,50)):
        self.signal = signal
        self.krange = krange

    def fcoef(self):
        # get Fourier coefficients of 1-dim function on [0,1]
        return fc(self.signal, krange=self.krange)

    def reconstruct(self, c=None):
        if c is None:
            c = self.fcoef()
        newsignal = np.real(fr(c, krange=self.krange, Nt=self.shape))
        return newsignal

    def filter(self, condition):
        # condition: the frequence will be filtered
        c = self.fcoef()
        for k, row in enumerate(c, self.krange[0]):
            if condition(k):
                c[k-self.krange[0]] = 0

        return self.reconstruct(c)


    def show(self, axes=None):
        if axes is None:
            plt.plot(self.signal)
        else:
            axes.plot(self.signal)


class ImageProcessor(SignalProcessor):
    # Image Processor with Fourier Analysis
    # krange: ft of image
    def __init__(self, krange=((-50,50), (-50,50))):
        self.krange = krange

    def fcoef(self, image):
        # get Fourier coefficients of 2-dim function representing an image
        image = np.float32(image)
        return [fc2(image[:,:,ch], image.shape[:2], self.krange[0], self.krange[1]) for ch in range(3)]

    def reconstruct(self, image, c=None):
        if c is None:
            c = self.fcoef(image)
        newimage = np.zeros(image.shape)
        newimage[:,:,0], newimage[:,:,1], newimage[:,:,2] = [np.real(fr2(c[i], self.krange[0], self.krange[1], image.shape[:2])[0]) for i in range(3)]
        convert = np.frompyfunc(lambda x: min(max(np.round(x), 0), 255), 1, 1)
        return np.uint8(convert(newimage))


    def filter(self, image, condition):
        # condition: the frequence will be filtered
        cs = self.fcoef(image)
        for c in cs:
            for k in range(self.krange[0][0], self.krange[0][1]):
                for j in range(self.krange[1][0], self.krange[1][1]):
                    if condition(k,j):
                        c[k-self.krange[0][0], j-self.krange[1][0]]=0

        return self.reconstruct(image, cs)

    def filterx(self, image, *condition):
        # condition: the frequence will be filtered
        cs = self.fcoef(image)
        for c in cs:
            for k in range(self.krange[0][0], self.krange[0][1]):
                if condition[0](k):
                    c[k-self.krange[0][0], :] = 0
                else:
                    for j in range(self.krange[1][0], self.krange[1][1]):
                        if condition[1](j):
                            c[k-self.krange[0][0], j-self.krange[1][0]]=0

        return self.reconstruct(image, cs)


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import numpy.linalg as LA

    def energy(x):
        return LA.norm(x)**2

    # Nt = 101
    # t = np.linspace(0,1,Nt)
    # f = np.sin(t*4)
    # n = f+ np.random.randn(Nt)/10
    #r, t = fapprox(n, threshold=0.1)
    # plt.plot(t, r, t,f, t, n, 'r')

    # image processing with Fourier analysis
    image = plt.imread('../src/Fourier.jpeg')
    ip = ImageProcessor()
    ax1 = plt.subplot(231)
    plt.title('Original Image')
    ax1.imshow(image)
    plt.axis('off')
    ax2 = plt.subplot(232)
    E = energy(image)
    high = ip.filterx(image, lambda x:-5<=x<5,lambda x:-5<=x<5)
    plt.title(f'High Frequence {energy(high)/E:.2%}')
    plt.imshow(255-high)
    plt.axis('off')
    plt.subplot(233)
    low = ip.filterx(image, lambda x:5<=x or x<-5,lambda x:5<=x or x<-5)
    plt.title(f'Low Frequence{energy(low)/E:.2%}')
    plt.imshow(low)
    plt.axis('off')
    plt.subplot(234)
    plt.title('Reconstruction')
    plt.imshow(ip.reconstruct(image))
    plt.axis('off')

    plt.subplot(235)
    plt.title('Almost Total Reconstruction')
    ip.krange = ((-100,100),(-150,150))
    plt.imshow(ip.reconstruct(image))
    plt.axis('off')
    
    # m = n = 100
    # r, c, _ = image.shape
    # I = np.float32(image)
    # A = np.array([np.kron(I[::r//m, ::c//n, 0], np.ones(r//m, c//n)),
    # np.kron(I[::r//m, ::c//n, 1], np.ones(r//m, c//n)),
    # np.kron(I[::r//m, ::c//n, 2], np.ones(r//m, c//n))])
    # A = A.transpose((1, 2, 0))
    # plt.subplot(235)
    # plt.imshow(ip.reconstruct(A))
    # plt.axis('off')
    plt.show()

